#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2025/02/01 2:40
@Author  : thezehui@gmail.com
@File    : chat.py
"""
import os
from typing import Tuple, List, Any, Dict, Union

import tiktoken
from langchain_core.messages import BaseMessage, get_buffer_string
from langchain_openai.chat_models.base import BaseChatOpenAI

from internal.core.language_model.entities.model_entity import BaseLanguageModel


# 从函数导入，而不是导入整个模块


class Chat(BaseChatOpenAI, BaseLanguageModel):
    """深度求索大语言模型基类"""

    def __init__(self, *args, **kwargs):
        super().__init__(
            *args,
            openai_api_key=os.getenv("DEEPSEEK_API_KEY"),
            openai_api_base=os.getenv("DEEPSEEK_API_BASE"),
            **kwargs
        )

    def _get_encoding_model(self) -> Tuple[str, tiktoken.Encoding]:
        """重写月之暗面获取编码模型名字+模型函数，该类继承OpenAI，词表模型可以使用gpt-3.5-turbo防止出错"""
        # 1.将DeepSeek的词表模型设置为gpt-3.5-turbo
        model = "deepseek-chat"

        # 2.返回模型名字+编码器
        # 由于tiktoken无法自动映射deepseek-chat模型，我们显式使用cl100k_base编码器
        # cl100k_base是GPT-4和GPT-3.5-Turbo使用的编码器，可以作为替代
        return model, tiktoken.get_encoding("cl100k_base")

    def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
        """重写获取消息token数量的方法，使用cl100k_base编码器计算token数量
        
        Args:
            messages: 消息列表
            
        Returns:
            消息的token数量
        """
        # 获取编码器
        _, encoding = self._get_encoding_model()

        # 计算所有消息的token数量
        num_tokens = 0
        for message in messages:
            # 获取消息的文本内容
            message_text = get_buffer_string([message])
            # 使用编码器计算token数量
            num_tokens += len(encoding.encode(message_text))

            # 如果消息有工具调用，也需要计算工具调用的token数量
            if hasattr(message, "tool_calls") and message.tool_calls:
                for tool_call in message.tool_calls:
                    # 工具名称
                    if "name" in tool_call:
                        num_tokens += len(encoding.encode(tool_call["name"]))
                    # 工具参数 (转为字符串计算)
                    if "args" in tool_call:
                        args_str = str(tool_call["args"])
                        num_tokens += len(encoding.encode(args_str))

        return num_tokens
